/**----------------------------------------------------------------------------
 * WinDigApp.cpp
 *-----------------------------------------------------------------------------
 * 
 *-----------------------------------------------------------------------------
 * All rights reserved by Noh,Yonghwan (fixbrain@gmail.com, unsorted@msn.com)
 *-----------------------------------------------------------------------------
 * 26:8:2011   14:50 created
**---------------------------------------------------------------------------*/
#include "stdafx.h"

#include "Win32Utils.h"
#include "ShareHeader.h"
#include "WDGIoData.h"
#include "WDGIoCode.h"

#include "ServiceManager.h"
#include "DebugMessage.h"






#define ARGS_LEFT(x) (argc-(x))

#define 	cExit				0
#define 	cStartDriver		1
#define 	cStopDriver			2
#define		cStartProtector		3
#define		cStopProtector		4

//
// internal functions
//
void WDGShowHelp(void);
int WDGShowCommand(DWORD* Cmd);
int WDGProcessCommand(IN SVCMGR_CTX ctx, IN DWORD Command);


/**----------------------------------------------------------------------------
    \brief  
    
    \param  
    \return
    \code
    
    \endcode        
-----------------------------------------------------------------------------*/
int _tmain(int argc, _TCHAR* argv[])
{
	int ret = -1;
	
	_ftprintf(stderr, _T("==================================================\n"));
	_ftprintf(stderr, _T(" WinDig (Windows Diagnostic Tool)\n"));
	_ftprintf(stderr, _T(" copyright(c) 2007, somma (fixbrain@gmail.com)\n"));
	_ftprintf(stderr, _T("==================================================\n"));


	WCHAR CurPath[MAX_PATH + 1] = {0};
	WCHAR DriverPath[MAX_PATH + 1]={0};
	GetCurrentDirectoryW(MAX_PATH, CurPath);
	if (TRUE != SUCCEEDED(StringCbPrintfW(
                            DriverPath, 
							sizeof(DriverPath), 
							L"%s\\%s.sys", 
							CurPath,
							DRIVERNAMEU
							)))
    {
        fwprintf(stderr, L"[ERR ] can not make driver path.");
        goto terminate;
    }

    // #1 install driver
    // 
    SVCMGR_CTX ctx = InstallDriver(DriverPath, DRIVERNAMEU, SERVICE_DISPLAY_NAME);
    if (NULL == ctx)
    {
        fwprintf(stderr, L"[ERR ] InstallDriver(driver path=%s) failed", DriverPath);
        goto terminate;
    }

	// #2 command loop
	//
	DWORD Cmd;
	for(;;)
	{
		if (0 != WDGShowCommand(&Cmd) )	
		{ 
            fwprintf(stderr, L"[ERR ] WDGShowCommand() failed\n");
            goto terminate;
		}		

        if (cExit == Cmd) 
        {
            fwprintf(stdout, L"[INFO] Exit command seleted....\n");
            goto terminate;
        }

		ret = WDGProcessCommand(ctx, Cmd);
        if (0 != ret)
        {
            fwprintf(stderr, L"[ERR ] WDGProcessCommand(Command=0x%08x) failed\n", Cmd);
			goto terminate;
        }
	}// for

	ret = 0;	// for success	

terminate:
    if (NULL != ctx){ReleaseContext(ctx);}

	return ret;
}

/**	---------------------------------------------------------------------------
	\brief	

	\param	
	\return			
	\code
	\endcode		
-----------------------------------------------------------------------------*/
int WDGShowCommand(DWORD* Cmd)
{
	_ftprintf(stdout, _T("\ncommand :\n"));
	_ftprintf(stdout, _T("   0 - exit \n"));
	_ftprintf(stdout, _T("   1 - start driver\n"));
	_ftprintf(stdout, _T("   2 - stop driver\n"));
	_ftprintf(stdout, _T("   3 - start protector\n"));
	_ftprintf(stdout, _T("   4 - stop protector\n"));
	_ftprintf(stdout, _T("============================================\n"));	
	_ftprintf(stdout, _T(">> "));	

	DWORD dv=0;
	std::cin >> dv;
	*Cmd = dv;
	return 0;
}



/**----------------------------------------------------------------------------
    \brief  
    
    \param  
    \return
    \code
    
    \endcode        
-----------------------------------------------------------------------------*/
int WDGProcessCommand(IN SVCMGR_CTX ctx, IN DWORD Command)
{
    _ASSERTE(NULL != ctx);
    if (NULL == ctx) return -1;

    DTSTATUS status=DTS_WINAPI_FAILED;
    int ret = -1;
    switch(Command)
    {
    case cStartDriver:
        {
            status = StartDriverService(ctx);
            if (TRUE != DT_SUCCEEDED(status))
            {
                fwprintf(stderr, L"[ERR ] StartDriverService() failed, status=0x%08x", status);
                ret = -1;
            }
            else
            {
                fwprintf(stdout, L"[INFO] driver service started successfully.");
                ret = 0;
            }
            break;
        }
    case cStopDriver:
        {
            status = StopDriverService(ctx);
            if (TRUE != DT_SUCCEEDED(status))
            {
                fwprintf(stderr, L"[ERR ] StopDriverService() failed, status=0x%08x", status);
                ret = -1;
            }
            else
            {
                fwprintf(stdout, L"[INFO] driver service stopped successfully.");
                ret = 0;
            }
            break;
        }
	case cStartProtector:
		{
			fwprintf(stdout, L"pid to protect : ");
			DWORD pid = 0;
			std::cin >> pid;
			
			WDG_START_PROTECTOR_INFO param;
			param.PidToProtect = pid;
			__asm int 3
			// install and run hook driver
			// 
			DWORD cbReturn=0;
			status = SendDriverCommand(
						ctx, 
						WDG_START_PROTECTOR, 
						(LPVOID)&param, 
						sizeof(WDG_START_PROTECTOR_INFO), 
						NULL, 
						0, 
						&cbReturn, 
						NULL
						);
			if (TRUE != DT_SUCCEEDED(status))
			{
				fwprintf(stderr, L"[ERR ] WDG_START_PROTECTOR failed, status=0x%08x", status);
				ret = -1;
			}
			else
			{
				fwprintf(stdout, L"[INFO] protector started successfully.");
				ret = 0;
			}
			break;
		}
	case cStopProtector:
		{
			// stop / uninstall hook driver
			// 
			DWORD cbReturn=0;
			status = SendDriverCommand(
						ctx, 
						WDG_STOP_PROTECTOR, 
						NULL, 
						0, 
						NULL, 
						0, 
						&cbReturn, 
						NULL
						);
			if (TRUE != DT_SUCCEEDED(status))
			{
				fwprintf(stderr, L"[ERR ] WDG_STOP_PROTECTOR failed, status=0x%08x", status);
				ret = -1;
			}
			else
			{
				fwprintf(stdout, L"[INFO] protector stopped successfully.");
				ret = 0;
			}
			break;
		}
    default:
        {
            return -1;
        }
    }

    return ret;
}
